import sys, os
from torch.utils.data import DataLoader
import torch

from tf_gridnet import TF_GridNet

sys.path.append(os.path.dirname(os.path.dirname(os.getcwd())))
from data_loader import getTestSet
from data_loader import AudioData
from si_snr_se_loss import loss
import config
import argparse
import tqdm


def main(args):
    ## 1. prepare model 
    ckpt = torch.load(args.model_path)
    model = TF_GridNet(B = args.block_num, device = args.device)
    model.load_state_dict(ckpt['model_state_dict'])
    model.to("cuda:"+str(args.device))
    model.eval()

    ## 2. prepare dataset
    x_test, y_test = getTestSet(args.test_path)
    dataset = AudioData(x_test,y_test, device = "cuda:"+str(args.device))
    ## 3. for loop, calculate si-snr, 

    total_loss = 0 
    ct = 0 
    pbar = tqdm.tqdm(dataset)
    with torch.no_grad():
        for x,y in pbar:
            x,y = x.to(args.device), y.to(args.device)
            x = x.unsqueeze(0)
            y = y.unsqueeze(0)
            output = model(x)
            l = loss(output, y, name="si_snr_se", device = args.device)
            total_loss +=l.item()
            ct+=1
            pbar.set_description(f"loss {total_loss / ct}")
    print(f"average si_snr loss {total_loss / ct}")
    pass

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_path", type = str, required = True)
    parser.add_argument("--block_num", type = int, required = True)
    parser.add_argument("--device",type = int, required = True)
    parser.add_argument("--test_path", type = str, required= True)
    args = parser.parse_args()
    main(args)

    pass